import requests
import pandas as pd
from datetime import datetime
import os

def get_zt_data():
    """
    获取涨停板信息分析数据
    """
    # 获取当前日期，格式为YYYYMMDD
    today = datetime.now().strftime("%Y%m%d")
    
    # API URL
    url = f"http://127.0.0.1:6666/api/public/stock_zt_pool_em?date={today}"
    
    try:
        # 发送GET请求
        print(f"正在请求API: {url}")
        response = requests.get(url)
        
        # 检查响应状态
        if response.status_code != 200:
            print(f"API请求失败，状态码: {response.status_code}")
            return None
        
        # 解析JSON响应
        data = response.json()
        
        # 检查数据结构
        if not isinstance(data, list):
            print("API返回的数据格式不正确，预期为列表")
            return None
        
        if not data:
            print("API返回的数据为空")
            return None
        
        print(f"成功获取到 {len(data)} 条数据")
        return data
        
    except Exception as e:
        print(f"请求API时发生错误: {e}")
        return None

def process_data(data):
    """
    处理数据，只保留特定字段并格式化
    """
    if not data:
        return None

    # 创建一个新的列表来存储处理后的数据
    processed_data = []

    # 字段映射（API返回的字段名 -> 我们想要的字段名）
    field_mapping = {
        "序号": "序号",
        "代码": "代码",
        "名称": "名称",
        "涨跌幅": "涨跌幅",
        "最新价": "最新价",
        "成交额": "成交额",
        "换手率": "换手率",
        "封板资金": "封板资金",
        "首次封板时间": "首次封板时间",
        "最后封板时间": "最后封板时间",
        "炸板次数": "炸板次数",
        "涨停统计": "涨停统计",
        "连板数": "连板数",
        "所属行业": "所属行业"
    }

    # 我们只保留这些字段
    fields_to_keep = ["序号", "名称", "涨跌幅", "最新价", "成交额", "换手率", "封板资金", "涨停统计", "连板数", "所属行业"]

    for item in data:
        processed_item = {}

        for api_field, our_field in field_mapping.items():
            if api_field in item and our_field in fields_to_keep:
                value = item[api_field]

                # 格式化特定字段
                if our_field == "涨跌幅":
                    # 涨跌幅保留两位小数
                    try:
                        value = round(float(value), 2)
                    except (ValueError, TypeError):
                        pass

                elif our_field == "成交额":
                    # 成交额转为亿作为单位，保留两位小数
                    try:
                        value = round(float(value) / 100000000, 2)  # 假设原单位为元，转换为亿
                    except (ValueError, TypeError):
                        pass

                elif our_field == "换手率":
                    # 换手率保留两位小数
                    try:
                        value = round(float(value), 2)
                    except (ValueError, TypeError):
                        pass

                elif our_field == "封板资金":
                    # 封板资金转为亿作为单位，保留两位小数
                    try:
                        value = round(float(value) / 100000000, 2)  # 假设原单位为元，转换为亿
                    except (ValueError, TypeError):
                        pass

                processed_item[our_field] = value

        processed_data.append(processed_item)

    return processed_data

def save_to_excel(data, output_file=None):
    """
    将数据保存为Excel文件
    """
    if not data:
        print("没有数据可保存")
        return None

    # 创建DataFrame并保存为Excel
    dateStamp = datetime.now().strftime("%Y%m%d")
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir = f"../spiltExcel/input/{dateStamp}"
    os.makedirs(output_dir, exist_ok=True)

    # 如果没有指定输出文件，则使用当前日期作为文件名
    if not output_file:
        today = datetime.now().strftime("%Y%m%d")
        output_file = f"{output_dir}/ztbfx_{timestamp}.xlsx"

    try:
        # 创建DataFrame
        df = pd.DataFrame(data)

        # 确保列的顺序
        columns = ["序号", "名称", "涨跌幅", "最新价", "成交额", "换手率", "封板资金", "涨停统计", "连板数", "所属行业"]
        df = df.reindex(columns=[col for col in columns if col in df.columns])

        # 保存为Excel
        df.to_excel(output_file, index=False)

        print(f"数据已成功保存到: {os.path.abspath(output_file)}")
        return output_file

    except Exception as e:
        print(f"保存Excel时发生错误: {e}")
        return None

def main():
    """
    主函数
    """
    print("=" * 50)
    print("涨停板信息分析工具")
    print("=" * 50)

    # 获取数据
    data = get_zt_data()

    if not data:
        print("获取数据失败，程序退出")
        return

    # 处理数据
    processed_data = process_data(data)

    if not processed_data:
        print("处理数据失败，程序退出")
        return

    # 保存为Excel
    output_file = save_to_excel(processed_data)

    if output_file:
        print("\n处理完成！")
        print(f"共处理 {len(processed_data)} 条数据")
        print(f"输出文件: {output_file}")
    else:
        print("保存Excel失败，程序退出")

def process():
    main()

if __name__ == "__main__":
    main()
